#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
court.py 使用案例
展示如何使用羽毛球场地检测功能
"""
import os
import cv2
import numpy as np
from matplotlib import pyplot as plt
from court import (
    Get_Lines,
    horizontal_vertical,
    clear_duplicate_lines,
    find_course,
    draw_course,
    cross_point
)
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']  # 支持中文显示
plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号


def detect_badminton_court(image_path):
    """
    检测羽毛球场地的主函数

    Args:
        image_path: 图像文件路径

    Returns:
        dict: 包含检测结果的字典
    """
    # 1. 读取图像
    image = cv2.imread(image_path)
    if image is None:
        print(f"❌ 无法读取图像: {image_path}")
        return None

    image_copy = image.copy()
    image_copy1 = image.copy()

    print("🔍 开始检测羽毛球场地...")

    # 2. 获取线条
    lines = Get_Lines(image_copy)
    if not lines:
        print("❌ 未检测到任何线条")
        return None

    print(f"✅ 检测到 {len(lines)} 条线条")

    # 3. 分类线条（横线和竖线）
    court_exist, h_lines, s_lines = horizontal_vertical(lines)

    if not court_exist:
        print("❌ 未找到足够的横线或竖线")
        return None

    print(f"✅ 横线数量: {len(h_lines)}")
    print(f"✅ 竖线数量: {len(s_lines)}")

    # 4. 清理重复线条
    clear_duplicate_lines(h_lines)

    # 5. 合并所有线条
    lines = h_lines + s_lines

    # 6. 寻找场地边界
    left_line, right_line, top_line, bottom_line = find_course(lines)

    if not left_line or not right_line or not top_line or not bottom_line:
        print("❌ 未找到完整的场地边界")
        return None

    print("✅ 成功找到场地边界")

    # 7. 绘制场地边界
    result_image = draw_course(image_copy, left_line, right_line, top_line, bottom_line)

    return {
        'original_image': image,
        'result_image': result_image,
        'lines': lines,
        'h_lines': h_lines,
        's_lines': s_lines,
        'court_lines': {
            'left': left_line,
            'right': right_line,
            'top': top_line,
            'bottom': bottom_line
        }
    }


def convert_person_coordinates(person_x, person_y, result):
    """
    将人在图像中的坐标转换为场地坐标

    Args:
        person_x, person_y: 人在图像中的坐标
        result: 检测结果字典

    Returns:
        tuple: (场地X坐标, 场地Y坐标)
    """
    left_line = result['court_lines']['left']
    right_line = result['court_lines']['right']
    top_line = result['court_lines']['top']
    bottom_line = result['court_lines']['bottom']

    # 计算左边界在person_y处的x坐标
    x1, y1, x2, y2, flag = left_line
    left_x = (person_y - y1) * (x2 - x1) / (y2 - y1) + x1

    # 计算右边界在person_y处的x坐标
    x1, y1, x2, y2, flag = right_line
    right_x = (person_y - y1) * (x2 - x1) / (y2 - y1) + x1

    # 计算像素到厘米的转换比例
    pixel_to_cm_x = 610 / (right_x - left_x)  # 场地宽度610cm

    # 计算X坐标（厘米）
    court_x = (person_x - left_x) * pixel_to_cm_x

    # 计算Y坐标（厘米）
    # 这里需要根据实际的顶线和底线来计算
    abs_dd = 132  # 这是一个示例值，实际应该根据检测结果计算
    pixel_to_cm_y = 410 / abs_dd  # 场地长度410cm
    court_y = ((top_line[1] + top_line[3]) / 2 - person_y) * pixel_to_cm_y

    return court_x, court_y


def visualize_results(result):
    """
    可视化检测结果

    Args:
        result: 检测结果字典
    """
    if result is None:
        return

    # 创建子图
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))

    # 原始图像
    axes[0, 0].imshow(cv2.cvtColor(result['original_image'], cv2.COLOR_BGR2RGB))
    axes[0, 0].set_title("原始图像", fontsize=14)
    axes[0, 0].axis('off')

    # 检测到的所有线条
    lines_image = result['original_image'].copy()
    for line in result['lines']:
        x1, y1, x2, y2, flag = line[0]
        if flag == 'h':
            cv2.line(lines_image, (x1, y1), (x2, y2), (0, 0, 255), 2)  # 红色横线
        elif flag == 's':
            cv2.line(lines_image, (x1, y1), (x2, y2), (0, 255, 0), 2)  # 绿色竖线

    axes[0, 1].imshow(cv2.cvtColor(lines_image, cv2.COLOR_BGR2RGB))
    axes[0, 1].set_title("检测到的线条", fontsize=14)
    axes[0, 1].axis('off')

    # 场地边界
    axes[1, 0].imshow(cv2.cvtColor(result['result_image'], cv2.COLOR_BGR2RGB))
    axes[1, 0].set_title("场地边界", fontsize=14)
    axes[1, 0].axis('off')

    # 坐标转换示例
    coord_image = result['original_image'].copy()
    person_x, person_y = 400, 300  # 示例人物坐标
    court_x, court_y = convert_person_coordinates(person_x, person_y, result)

    # 标记人物位置
    cv2.circle(coord_image, (person_x, person_y), 10, (255, 0, 0), -1)
    cv2.putText(coord_image, f"Person: ({person_x}, {person_y})",
                (person_x + 15, person_y - 15), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
    cv2.putText(coord_image, f"Court: ({court_x:.1f}, {court_y:.1f}) cm",
                (person_x + 15, person_y + 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

    axes[1, 1].imshow(cv2.cvtColor(coord_image, cv2.COLOR_BGR2RGB))
    axes[1, 1].set_title("坐标转换示例", fontsize=14)
    axes[1, 1].axis('off')

    plt.tight_layout()
    plt.show()

def detect_badminton_court_from_frame(frame):
    """
    从视频帧检测羽毛球场地
    
    Args:
        frame: 视频帧图像
    
    Returns:
        dict: 检测结果
    """
    frame_copy = frame.copy()

    height, width = frame_copy.shape[:2]
    print(f"🔍 帧尺寸: {width}x{height}")
    
    # 获取线条
    lines = Get_Lines(frame_copy)
    if not lines:
        return None
    
    # 分类线条
    court_exist, h_lines, s_lines = horizontal_vertical(lines)
    if not court_exist:
        return None
    
    # 清理重复线条
    clear_duplicate_lines(h_lines)
    
    # 合并线条
    lines = h_lines + s_lines
    
    # 寻找场地边界
    left_line, right_line, top_line, bottom_line = find_course(lines)
    
    if not left_line or not right_line or not top_line or not bottom_line:
        return None
    
    # 绘制场地边界
    result_image = draw_course(frame_copy, left_line, right_line, top_line, bottom_line)
    
    return {
        'original_image': frame,
        'result_image': result_image,
        'lines': lines,
        'h_lines': h_lines,
        's_lines': s_lines,
        'court_lines': {
            'left': left_line,
            'right': right_line,
            'top': top_line,
            'bottom': bottom_line
        }
    }

# def main():
#     """
#     主函数 - 使用示例
#     """
#     print("🏸 羽毛球场地检测系统使用示例")
#     print("=" * 50)

#     # 示例1: 使用指定图像文件
#     image_path = "image/IMG20220815235720.jpg"  # 替换为你的图像路径

#     print(f"📁 正在处理图像: {image_path}")

#     # 检测场地
#     result = detect_badminton_court(image_path)

#     if result:
#         print("\n✅ 检测成功!")
#         print(f"检测到 {len(result['lines'])} 条线条")
#         print(f"横线: {len(result['h_lines'])} 条")
#         print(f"竖线: {len(result['s_lines'])} 条")

#         # 坐标转换示例
#         person_x, person_y = 0, 0  # 假设人在图像中的位置
#         court_x, court_y = convert_person_coordinates(person_x, person_y, result)
#         print(f"\n📍 坐标转换示例:")
#         print(f"图像坐标: ({person_x}, {person_y})")
#         print(f"场地坐标: ({court_x:.2f}, {court_y:.2f}) cm")

#         # 可视化结果
#         print("\n🖼️ 显示检测结果...")
#         visualize_results(result)

#     else:
#         print("❌ 检测失败，请检查图像质量或路径")

def draw_upper_filter_line(image, upper_ratio=0.5, line_color=(0, 255, 255), line_thickness=2):
    """
    在图像上画一条线标记上半部分过滤区域
    
    Args:
        image: 输入图像
        upper_ratio: 上半部分比例，默认0.5（50%）
        line_color: 线条颜色，默认黄色 (BGR格式)
        line_thickness: 线条粗细
    
    Returns:
        numpy.ndarray: 画了标记线的图像
    """
    image_copy = image.copy()
    height, width = image_copy.shape[:2]

    print(height)
    print(height * upper_ratio)

    # 计算过滤线位置
    filter_y = int(height * upper_ratio)
    
    # 画水平线标记过滤区域
    cv2.line(image_copy, (0, filter_y), (width, filter_y), line_color, line_thickness)
    
    # 添加文字说明
    text = f"Filter Line (Top {upper_ratio*100:.0f}%)"
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.8
    text_thickness = 2
    
    # 计算文字位置
    text_size = cv2.getTextSize(text, font, font_scale, text_thickness)[0]
    text_x = 10
    text_y = filter_y - 10
    
    # 画文字背景
    cv2.rectangle(image_copy, 
                  (text_x - 5, text_y - text_size[1] - 5), 
                  (text_x + text_size[0] + 5, text_y + 5), 
                  (0, 0, 0), -1)
    
    # 画文字
    cv2.putText(image_copy, text, (text_x, text_y), font, font_scale, line_color, text_thickness)
    
    return image_copy


def main():
    """
    主函数 - 使用示例
    """
    print("🏸 羽毛球场地检测系统使用示例")
    print("=" * 50)
    
    # 选择处理模式
    mode = input("选择处理模式 (1-图片, 2-视频, 3-摄像头): ").strip()
    
    if mode == "1":
        # 图片处理模式
        image_path = input("请输入图片路径 (回车使用默认): ").strip()
        if not image_path:
            image_path = "output/frame_000001_result.jpg"
        
        print(f"📁 正在处理图像: {image_path}")

        image = cv2.imread(image_path)
        if image is None:
            print(f"❌ 无法读取图像: {image_path}")
            return None

        # 检测场地
        result = detect_badminton_court_from_frame(image)

        if result:
            print("\n✅ 检测成功!")
            print(f"检测到 {len(result['lines'])} 条线条")
            print(f"横线: {len(result['h_lines'])} 条")
            print(f"竖线: {len(result['s_lines'])} 条")
            
            # 坐标转换示例
            person_x, person_y = 0, 0  # 假设人在图像中的位置
            court_x, court_y = convert_person_coordinates(person_x, person_y, result)
            print(f"\n📍 坐标转换示例:")
            print(f"图像坐标: ({person_x}, {person_y})")
            print(f"场地坐标: ({court_x:.2f}, {court_y:.2f}) cm")
            
            # 可视化结果
            print("\n🖼️ 显示检测结果...")
            visualize_results(result)
        else:
            print("❌ 检测失败，请检查图像质量或路径")
    
    elif mode == "2":
        # 视频处理模式
        video_path = input("请输入视频路径: ").strip()

        if not video_path:
            video_path = "image/202510051934.mp4"

        if video_path is None:
            print(f"❌ 无法读取图像: {video_path}")
            return None

        print(f"📹 正在处理视频: {video_path}")
        
        # 打开视频
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            print(f"❌ 无法打开视频文件: {video_path}")
            return
        
        # 获取视频信息
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        
        print(f"📹 视频信息: {width}x{height}, {fps}FPS, {total_frames}帧")
        
        # 处理视频
        frame_count = 0
        success_count = 0
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            frame_count += 1
            
            # 处理当前帧
            result = detect_badminton_court_from_frame(frame)

            if result:
                success_count += 1

                result['result_image'] = draw_upper_filter_line(result['result_image']);

                # 显示处理后的帧
                cv2.imshow('Video Processing', result['result_image'])
                # 修复保存问题
                output_dir = 'output'
                output_dir_end = 'output_end'
                if not os.path.exists(output_dir):
                    os.makedirs(output_dir)

                if not os.path.exists(output_dir_end):
                    os.makedirs(output_dir_end)

                # 保存处理后的图像，而不是原始帧
                output_path = f'{output_dir}/frame_{frame_count:06d}_result.jpg'
                output_path_end = f'{output_dir_end}/frame_{frame_count:06d}_result_end.jpg'
                cv2.imwrite(output_path,frame)
                cv2.imwrite(output_path_end, result['result_image'])
                print(f"💾 保存图像: {output_path}")
            else:
                frame = draw_upper_filter_line(frame)
                # 显示原始帧
                cv2.imshow('Video Processing', frame)

            
            # 显示进度
            if frame_count % (fps * 2) == 0:  # 每2秒显示一次
                progress = (frame_count / total_frames) * 100
                print(f"进度: {progress:.1f}% ({frame_count}/{total_frames})")
            
            # 按 'q' 退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        
        # 清理资源
        cap.release()
        cv2.destroyAllWindows()
        
        print(f"\n📊 处理完成!")
        print(f"总帧数: {frame_count}")
        print(f"成功检测: {success_count}")
        print(f"成功率: {(success_count/frame_count*100):.1f}%")
    
    elif mode == "3":
        # 摄像头处理模式
        camera_id = input("请输入摄像头ID (回车使用默认0): ").strip()
        if not camera_id:
            camera_id = 0
        else:
            camera_id = int(camera_id)
        
        print(f"📹 正在打开摄像头 {camera_id}...")
        
        # 打开摄像头
        cap = cv2.VideoCapture(camera_id)
        if not cap.isOpened():
            print(f"❌ 无法打开摄像头 {camera_id}")
            return
        
        # 设置摄像头参数
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
        cap.set(cv2.CAP_PROP_FPS, 30)
        
        # 获取实际参数
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        
        print(f"✅ 摄像头初始化成功")
        print(f"   分辨率: {width}x{height}")
        print(f"   帧率: {fps} FPS")
        print("按 'q' 退出, 按 's' 保存当前帧")
        
        # 处理摄像头
        frame_count = 0
        success_count = 0
        start_time = time.time()
        
        while True:
            ret, frame = cap.read()
            if not ret:
                print("❌ 无法读取摄像头帧")
                break
            
            frame_count += 1
            
            # 处理当前帧
            result = detect_badminton_court_from_frame(frame)

            if result:
                success_count += 1
                result['result_image'] = draw_upper_filter_line(result['result_image'])
                # 显示处理后的帧
                cv2.imshow('Camera Processing', result['result_image'])
            else:
                frame = draw_upper_filter_line(frame)
                # 显示原始帧
                cv2.imshow('Camera Processing', frame)
            
            # 显示统计信息
            if frame_count % 30 == 0:  # 每30帧显示一次
                elapsed_time = time.time() - start_time
                current_fps = frame_count / elapsed_time if elapsed_time > 0 else 0
                success_rate = (success_count / frame_count) * 100 if frame_count > 0 else 0
                print(f"帧数: {frame_count}, 成功: {success_count}, 成功率: {success_rate:.1f}%, FPS: {current_fps:.1f}")
            
            # 处理按键
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                print("👋 退出摄像头检测")
                break
            elif key == ord('s'):
                # 保存当前帧
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                if result:
                    save_path = f"camera_success_{timestamp}.jpg"
                    cv2.imwrite(save_path, result['result_image'])
                    print(f"💾 保存成功检测: {save_path}")
                else:
                    save_path = f"camera_failed_{timestamp}.jpg"
                    cv2.imwrite(save_path, frame)
                    print(f"💾 保存失败检测: {save_path}")
        
        # 清理资源
        cap.release()
        cv2.destroyAllWindows()
        
        # 显示最终统计
        elapsed_time = time.time() - start_time
        final_fps = frame_count / elapsed_time if elapsed_time > 0 else 0
        final_success_rate = (success_count / frame_count) * 100 if frame_count > 0 else 0
        
        print(f"\n📊 摄像头检测完成!")
        print(f"总帧数: {frame_count}")
        print(f"成功检测: {success_count}")
        print(f"成功率: {final_success_rate:.1f}%")
        print(f"平均FPS: {final_fps:.1f}")
        print(f"运行时间: {elapsed_time:.1f}秒")
    
    else:
        print("❌ 无效选择，请输入 1、2 或 3")

if __name__ == "__main__":
    # 运行示例
    main()

    # 你也可以直接调用检测函数
    # result = detect_badminton_court("your_image.jpg")
    # if result:
    #     visualize_results(result)